from __future__ import division
import numpy as np
#from numpy import multiply as nm
from Constant_Physics import Physics
import matplotlib.pyplot as plt


class Bfield():
    def __init__(self,B0=0.1,B0xyz=np.array([0, 0, 1])):
        self.B0=B0
        self.B0xyz=B0xyz
    
    def field_vec(self):
        res=np.array([0,0,0])
        res = (self.B0 ) / np.linalg.norm(self.B0xyz) * self.B0xyz    
        return res
        
    def field_norm(self):
        return self.B0


class Shallow_Donor(Physics):

    def __init__(self, clk, specie="P"):
        if specie == "P":
            self.nspin = 1 / 2
            self.dim_n=2
            self.A = 117.53    #MHz
            self.ge = 1.9985
            self.gn = -2.2632
        elif specie == "As":
            self.nspin = 3 / 2
            self.A = 198.35   #MHz
            self.ge = 1.9984
            self.gn = -0.762601
            self.dim_n=4
        elif specie == "Sb5/2":
            self.nspin = 5 / 2
            self.A = 186.802   #MHz
            self.ge = 1.9986
            self.gn = -1.3454
            self.dim_n=6
        elif specie == "Sb7/2":
            self.nspin = 7 / 2
            self.A = 101.516    #MHz
            self.ge = 1.9986
            self.gn = -0.728492
            self.dim_n=8
        elif specie == "Bi":
            self.nspin = 9 / 2
            self.A = 1475.4    #MHz
            self.ge = 2.0003
            self.gn = -0.914
            self.dim_n=10

        self.dim_e=2
        self.espin=1/2
        self.dimension = self.dim_e* self.dim_n
        self.hamiltonian = np.zeros([self.dimension, self.dimension])
        self.transformation_matrix = np.zeros([self.dimension, self.dimension])
        
        self.wn=Physics.EnergyUnit * (10 ** (-6)) / Physics.h * self.gn * Physics.mun
        self.we=Physics.EnergyUnit * (10 ** (-6)) / Physics.h * self.ge * Physics.mue


    
        #Preparation of the nuclear spin matrices
        self.Ip = np.zeros([self.nspin * 2 + 1, self.nspin * 2 + 1])
        list_jm = np.arange(-self.nspin, self.nspin + 1)
        for n in range(int(self.nspin * 2)):
            self.Ip[n , n + 1] = np.sqrt((self.nspin - list_jm[n]) * (self.nspin + list_jm[n] + 1))
        self.Ix = (1 / 2) * (self.Ip + np.transpose(self.Ip))
        self.Iy = (1 / 2) * (np.transpose(self.Ip) - self.Ip)
        self.Iz = np.diag(-list_jm)
        self.I = [self.Ix, self.Iy, self.Iz]
        self.Id = np.diag(np.ones(self.nspin * 2 + 1))

        self.espin = 1 / 2  #Electronic spin
        self.dim_e=2
        self.Sx = 1 / 2 * np.array([[0, 1], [ 1, 0]])
        self.Sy = 1 / 2 * np.array([[0, -1], [1, 0]])
        self.Sz = 1 / 2 * np.array([[1, 0], [0, -1]])
        self.S = [self.Sx, self.Sy, self.Sz]
        self.Sd = np.diag([1, 1])

    def Hamiltonian(self, b):
        """
        Input:
            -B: the magnetic field [Bx,By,Bz]
        Output:
            -hamiltonian: the donor hamiltonian (Ze+Zn+H) 
        """
        Ze=self.we*(self.Sx*b[0]+self.Sy*b[1]+self.Sz*b[2])
        Zn=self.wn*(self.Ix*b[0]+self.Iy*b[1]+self.Iz*b[2])
        hyp=self.A*(np.kron(self.Sx,self.Ix)-np.kron(self.Sy,self.Iy)+np.kron(self.Sz,self.Iz))
        return np.kron(Ze,self.Id) + np.kron(self.Sd,Zn)+ hyp
    
    def Eigenvalues(self,b):
        self.eigenval, self.eigenvec = np.linalg.eig(self.Hamiltonian(b))
        return self.eigenval
    
    def EPR_transitions(self,b):
        """
        This is true P donor EPR transitions until 500 gauss. 
        Then, the calculation is false (fucking eigensystem that messes up the order of the eigenvalues...)
        """
        self.eigenval, self.eigenvec = np.linalg.eig(self.Hamiltonian(b))
        n=int(self.nspin*2+1)
        epr=np.zeros(n)
        for i in range(n):
            epr[i]=np.abs(self.eigenval[i])+np.abs(self.eigenval[i+n])
        return epr


if __name__=="__main__":
    donor=Shallow_Donor("P")
    b=Bfield(B0=0.05)
    print donor.Eigenvalues(b.field_vec())
    print donor.EPR_transitions(b.field_vec())